#!/usr/bin/env python3
"""修复ChromaDB路径配置问题"""

import os
import shutil
import asyncio
from pathlib import Path
from loguru import logger

# 添加src目录到Python路径
import sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from config import load_config
from database.manager import DatabaseManager

async def fix_chromadb_path_configuration():
    """修复ChromaDB路径配置问题"""
    
    # 1. 加载配置
    config = load_config()
    chromadb_config = config.get('databases', {}).get('chromadb', {})
    configured_path = chromadb_config.get('persist_directory', './data/chroma_db')
    
    print(f"配置文件中的ChromaDB路径: {configured_path}")
    
    # 2. 检查现有的ChromaDB目录
    project_root = Path(__file__).parent
    paths_to_check = [
        project_root / 'chroma_db',
        project_root / 'data' / 'chroma_db', 
        project_root / 'data' / 'chromadb',
        project_root / configured_path.lstrip('./')
    ]
    
    existing_paths = []
    for path in paths_to_check:
        if path.exists() and path.is_dir():
            file_count = len(list(path.rglob('*')))
            size = sum(f.stat().st_size for f in path.rglob('*') if f.is_file())
            existing_paths.append({
                'path': path,
                'file_count': file_count,
                'size_mb': size / (1024 * 1024)
            })
            print(f"发现ChromaDB目录: {path} (文件数: {file_count}, 大小: {size / (1024 * 1024):.2f}MB)")
    
    if not existing_paths:
        print("未发现任何ChromaDB目录")
        return
    
    # 3. 选择最大的目录作为主目录
    main_path_info = max(existing_paths, key=lambda x: x['size_mb'])
    main_path = main_path_info['path']
    configured_full_path = project_root / configured_path.lstrip('./')
    
    print(f"\n选择主ChromaDB目录: {main_path} (大小: {main_path_info['size_mb']:.2f}MB)")
    print(f"配置文件指定路径: {configured_full_path}")
    
    # 4. 如果主目录不是配置文件指定的路径，则移动数据
    if main_path != configured_full_path:
        print(f"\n需要将数据从 {main_path} 移动到 {configured_full_path}")
        
        # 确保目标目录的父目录存在
        configured_full_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 如果目标目录已存在，先备份
        if configured_full_path.exists():
            backup_path = configured_full_path.with_suffix('.backup')
            print(f"备份现有目录到: {backup_path}")
            if backup_path.exists():
                shutil.rmtree(backup_path)
            shutil.move(str(configured_full_path), str(backup_path))
        
        # 移动主目录到配置路径
        print(f"移动 {main_path} 到 {configured_full_path}")
        shutil.move(str(main_path), str(configured_full_path))
        
        # 清理其他ChromaDB目录
        for path_info in existing_paths:
            path = path_info['path']
            if path != main_path and path.exists():
                print(f"删除重复目录: {path}")
                shutil.rmtree(path)
    
    # 5. 验证修复结果
    print("\n验证ChromaDB配置...")
    
    try:
        # 初始化数据库管理器
        db_manager = DatabaseManager(config)
        await db_manager.initialize()
        
        # 检查ChromaDB连接
        if await db_manager.chromadb_client.ping():
            print("✓ ChromaDB连接成功")
            
            # 获取统计信息
            stats = await db_manager.chromadb_client.get_stats()
            print(f"✓ ChromaDB统计信息: {stats}")
            
        else:
            print("✗ ChromaDB连接失败")
        
        await db_manager.close()
        
    except Exception as e:
        print(f"✗ 验证失败: {e}")
    
    print("\nChromeDB路径配置修复完成！")

if __name__ == "__main__":
    asyncio.run(fix_chromadb_path_configuration())